# AUTHOR: DING
# -*- codeing = utf-8 -*-
# @Time: 2024/3/1 9:19
# @Author: 86139
# @Site: 
# @File: 22-train.py
# @Software: PyCharm
# tensorboard --logdir=pytorch/logs --port=6007

import torch
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
from model import *

train_data = torchvision.datasets.CIFAR10("./dataset", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
train_size = len(train_data)
test_size = len(test_data)
print("训练集的长度为:{}".format(train_size))
print("测试集的长度为:{}".format(test_size))
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# 网络模型
network = MyModule()
network.load_state_dict(torch.load("19-optim.pth"))
# 优化器
learning_rate = 1e-4
optim = torch.optim.SGD(network.parameters(), lr=learning_rate)
# 损失函数
loss = nn.CrossEntropyLoss()

# 训练参数
total_train_step = 0
total_test_step = 0
epoch = 3
# 添加tensorboard
writer = SummaryWriter("./logs")

for i in range(epoch):
    print("----------第{}轮开始----------".format(i + 1))
    # 开始训练
    network.train()
    t0 = time.time()
    for data in train_dataloader:
        imgs, targets = data
        outputs = network(imgs)
        # 损失
        output_loss = loss(outputs, targets)
        # 梯度下降
        optim.zero_grad()
        output_loss.backward()
        optim.step()
        if total_train_step % 100 == 0:
            t1 = time.time()
            print(t1 - t0)
            print("训练次数：{},loss:{}".format(total_train_step, output_loss.item()))
            writer.add_scalar("train_loss", output_loss.item(), total_train_step)
        total_train_step = total_train_step + 1

    # 开始测试
    network.eval()
    with torch.no_grad():
        total_test_loss = 0  # 测试集的loss和
        total_right_num = 0  # 整个测试集的准确率

        for data in test_dataloader:
            imgs, targets = data
            outputs = network(imgs)
            test_loss = loss(outputs, targets)
            total_test_loss = total_test_loss + test_loss
            right_num = (outputs.argmax(1) == targets).sum() # 计算判断正确的个数
            total_right_num = total_right_num + right_num
    print("整个测试集上的loss:{}".format(total_test_loss))
    print("整个测试集上的识别准确率:{}".format(total_right_num/test_size))
    writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
    writer.add_scalar("accuracy", total_right_num/test_size, total_test_step)
    total_test_step = total_test_step + 1

torch.save(network.state_dict(), "22-model.pth")

writer.close()
